from functools import partial
import math


def lr_schedule(decay_type, 
                learning_rate, 
                min_lr, 
                total_iters, 
                warmup_iters_ratio=0.05, 
                warmup_lr_ratio=0.1, 
                no_aug_iter_ratio=0.05, 
                step_num=10):
    
    def yolox_warm_cos_lr(learning_rate, 
                          min_lr, 
                          total_iters, 
                          warmup_iters, 
                          warmup_lr, 
                          no_aug_iter, 
                          iterations):
        iterations = int(iterations)
        if iterations <= warmup_iters:
            learning_rate = (learning_rate - warmup_lr) * pow(iterations/warmup_iters, 2) + warmup_lr
        elif iterations >= total_iters - float(no_aug_iter):
            learning_rate = min_lr
        else:
            learning_rate = min_lr + 0.5 * (learning_rate - min_lr) * (
                1.0 + math.cos(math.pi * ((iterations - warmup_iters)/(total_iters-warmup_iters-no_aug_iter))))
        return learning_rate

    def step_lr(learning_rate, 
                decay_rate, 
                step_size, 
                iterations):
        iterations = int(iterations)
        if step_size < 1:
            raise ValueError("step_size must above 1.")

        n       = iterations / step_size
        out_lr  = learning_rate * decay_rate ** n
        return out_lr

    if decay_type == "cos":
        warmup_iters = min(max(warmup_iters_ratio * total_iters, 1), 3)
        warmup_lr    = max(warmup_lr_ratio * learning_rate, 1e-6)
        no_aug_iter  = min(max(no_aug_iter_ratio * total_iters, 1), 15)
        lr_schedule_func = partial(yolox_warm_cos_lr, 
                                   learning_rate, 
                                   min_lr, 
                                   total_iters, 
                                   warmup_iters, 
                                   warmup_lr, 
                                   no_aug_iter)
    else:
        decay_rate = (min_lr / learning_rate) ** (1 / (step_num - 1))
        step_size  = total_iters / step_num
        lr_schedule_func = partial(step_lr, 
                                   learning_rate, 
                                   decay_rate, 
                                   step_size)

    return lr_schedule_func